﻿import os
import struct
import numpy as np
from PIL import Image
import gzip

def read_idx(filename):
    """读取 IDX 格式文件并返回 numpy 数组。"""
    with gzip.open(filename, 'rb') as f:  # 使用 gzip.open 读取压缩文件
        zero, data_type, dims = struct.unpack('>HBB', f.read(4))
        shape = tuple(struct.unpack('>I', f.read(4))[0] for _ in range(dims))
        return np.frombuffer(f.read(), dtype=np.uint8).reshape(shape)

def save_mnist_images(image_file, label_file, output_dir):
    """从 MNIST 数据集中生成并保存图片。

    Args:
        image_file (str): 图像数据文件路径。
        label_file (str): 标签数据文件路径。
        output_dir (str): 保存生成图片的目录。
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 读取图像和标签
    images = read_idx(image_file)
    labels = read_idx(label_file)

    #images[0][7][15] 表示第0张照片的第7行的第15个像素
    #labels[0] 表示第0个照片是什么数字

    # 遍历并保存每张图片
    # for i in range(images.shape[0]): # 这里如果6万张图片太难了，帮我选2个类别做区分把
    for i in range(1000):
        img = Image.fromarray(images[i], mode='L')  # 'L' 表示灰度图像
        label = labels[i]
        # 保存文件，以标签为前缀命名图片文件
        img.save(os.path.join(output_dir, f'{label}_{i}.png'))

# 获取当前文件所在目录
current_directory = os.path.dirname(os.path.abspath(__file__))

# 压缩文件路径
train_image_file = os.path.join(current_directory, 'File\\train-images-idx3-ubyte.gz')
train_label_file = os.path.join(current_directory, 'File\\train-labels-idx1-ubyte.gz')
test_image_file = os.path.join(current_directory, 'File\\t10k-images-idx3-ubyte.gz')
test_label_file = os.path.join(current_directory, 'File\\t10k-labels-idx1-ubyte.gz')

# 指定输出目录
train_output_dir = os.path.join(current_directory, 'Images\\mnist_train_images')
test_output_dir = os.path.join(current_directory, 'Images\\mnist_test_images')

# 生成并保存训练集和测试集图片
save_mnist_images(train_image_file, train_label_file, train_output_dir)
save_mnist_images(test_image_file, test_label_file, test_output_dir)

# 这里如果6万张图片太难了，帮我选2个类别做区分把
